package com.example.aigc_education.ws;

import com.example.aigc_education.utils.SecurityUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

import java.security.Principal;
import java.util.Map;

@Slf4j
public class SecurityHandshakeHandler extends DefaultHandshakeHandler {

    @Override
    protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler,
                                      Map<String, Object> attributes) {
        // 从attributes获取认证信息
        if (attributes.containsKey("SPRING_SECURITY_CONTEXT")) {
            Authentication authentication =
                    ((org.springframework.security.core.context.SecurityContext) attributes.get("SPRING_SECURITY_CONTEXT"))
                            .getAuthentication();
            if (authentication != null) {
                Long userId = SecurityUtils.getUserId();
                if (userId != null) {
                    log.info("WebSocket连接用户ID: {}", userId);
                    return authentication;
                }
                log.info("WebSocket连接用户: {}", authentication.getName());
                return authentication;
            }
        }

        // 从当前上下文获取认证信息
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        if (authentication != null) {
            log.info("从SecurityContext获取WebSocket连接用户: {}", authentication.getName());
            return authentication;
        }

        log.warn("未能确定WebSocket连接用户身份");
        return null;
    }
}
